from dataclasses import dataclass
from enum import Enum, auto
from typing import Optional, Tuple, Union

import numpy as np
import pydantic_numpy.typing as pnd
from pydantic import BaseModel, Field
from scipy.optimize import OptimizeResult


##0000000000000000000000000000000000000000000000000000000000000000000000000000000
## default constants
@dataclass
class DftValues:
    ## D -> Default; LM -> Lambda parameters; RLXD -> Relaxed
    fctr_symmetry = 2.0  # 对于 cylinder 探测器, Phi 方向没有分bin, R 方向 从 1 到 N 个 bin; 相当于高斯分布只有对称的一半
    epsilon = 1e-12  # Small value to prevent division by zero
    penalty_sigma_smooth = 1e3  # Lambda for sigma smooth penalty
    penalty_norm = 1e3  # Lambda for normalization penalty
    sigma_min = 1e-3  # Minimum allowed sigma value
    weight_min = 1e-8  # Minimum allowed weight
    norm_fctr_min = 0.5  # Minimum normalization factor
    norm_fctr_max = 1.5  # Maximum normalization factor
    step_range = 0.2  # Parameter step range for adjacent layer fitting
    step_rlxd_range = 5.0  # Relaxed step range for post-Bragg region
    smoothwindow_size = 15  # Window size for smoothing in mid-tail detection
    clip_minwidth = 2.0  # Minimum width for mid-tail detection
    clip_maxwidth = 40.0  # Maximum width for mid-tail detection
    lmbd_clip_min = 6.0  # Minimum lambda_m for mid-tail detection
    lmbd_clip_max = 40.0  # Maximum lambda_m for mid-tail detection
    rltv_err = 0.01  # Relative error threshold
    abs_err = 0.0  # Absolute error threshold
    gap_threshold = 0.15  # Gap threshold for continuous region detection (1.5 * dx)

    ## Default values for mid_tail_weights; mtl -> Mid Tail
    mtl_center = 15.0
    mtl_width = 3.0
    mtl_lmbd = 10.0

    ## Default ranges for mid_tail_detect
    mtl_min = 5.0
    mtl_rmax = 40.0
    clip_range = (clip_minwidth, clip_maxwidth)
    lmbds_clip = (lmbd_clip_min, lmbd_clip_max)


##xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx singleton
myDftVals = DftValues()


##0000000000000000000000000000000000000000000000000000 fitutils.py
##================================================= 网格类型
class MeshEnum(Enum):
    Box = auto()
    Cyliner = auto()


##================================= 归一化的输入
class NmlzIn(BaseModel):
    mtype: MeshEnum = Field(...)
    vars_tgt: pnd.Np1DArray = Field(...)  ## e.g. values target, 例如 dose
    bin_width: float = Field(...)


##================================= weighted std dev 的输入
class WtdStdDevIn(BaseModel):
    qty: pnd.Np1DArray = Field(...)
    wts: pnd.Np1DArray = Field(...)
    eps: float = myDftVals.epsilon


##xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx mid tail detect
class MidTailDtctData(BaseModel):
    crds: pnd.Np1DArray = Field(...)
    vals_tgt: pnd.Np1DArray = Field(...)  ## 测量值, 目标值
    vals_p0: pnd.Np1DArray = Field(...)  ## values of predict run 0


class MidTailDtctCfg(BaseModel):
    """
    Mid-tail detection parameters.
    """

    r_min: float = DftValues.mtl_min
    r_max: float = DftValues.mtl_rmax
    smooth_window: int = DftValues.smoothwindow_size
    width_clip: Tuple[float, float] = (DftValues.clip_minwidth, DftValues.clip_maxwidth)
    lam_m_clip: Tuple[float, float] = (DftValues.lmbd_clip_min, DftValues.lmbd_clip_max)
    rel_err: float = DftValues.rltv_err
    abs_err: float = DftValues.abs_err


class MidTailDcttRet(BaseModel):
    crds: pnd.Np1DArray = Field(...)  ## coordinates
    center: float
    width: float
    lam_m: float
    left: Optional[float]
    right: Optional[float]


## 归一化的输入
# @dataclass
# class MidTailWtsIn:
#     crds: pnd.Np1DArray
#     center: float = dftVals.mtl_center
#     width: float = dftVals.mtl_width
#     lam_m: float = dftVals.mtl_lmbd


##000000000000000000000000000000000000000000000000000000000000000 multigauss.py
##================================================== 多高斯模型, 参数转换工具
class GaussParams(BaseModel):
    """
    在fitting中, 高斯模型的独立参数(自由度) pack 到一个 array;
    此工具方便在 array 和 dataclass 之间转换.
    结构体字段具名， 方便访问
    """

    norm_factor: float  # Overall normalization factor
    weights: pnd.Np1DArray  # Weights [w2, w3, ..., wn] for components 2 to n (first weight is 1 - sum of others)
    sigmas: pnd.Np1DArray  # Sigmas [s1, s2, ..., sn] for each Gaussian component
    ngauss: int = 5  # 高斯分量个数
    with_background: bool = False  ## with back_ground
    background: Optional[float] = None  # 可选, background constant

    @property
    def w0(self) -> float:
        """Calculate the first component weight as 1 - sum of other weights."""
        return max(0.0, 1.0 - np.sum(self.weights))  # Ensure non-negative

    @property
    def all_weights(self) -> pnd.Np1DArray:
        """Compose all weights"""
        return np.concatenate(([self.w0], self.weights))

    @property
    def sum_weights2n(self) -> float:
        return np.sum(self.weights)

    @classmethod
    def from_array(cls, params: pnd.Np1DArray, n_gauss: int, with_background: bool = False):
        """
        raw parameter array -> GaussParams object
        """
        norm_factor = params[0]
        weights = params[1:n_gauss]  # weights for components 2 to n
        sigmas = params[n_gauss : 2 * n_gauss]  # sigmas for all components

        background = None
        if with_background:
            background = params[-1]

        return cls(
            norm_factor=norm_factor,
            weights=weights,
            sigmas=sigmas,
            ngauss=n_gauss,
            with_background=with_background,
            background=background,
        )

    def to_array(self) -> pnd.Np1DArray:
        """
        GaussParams object -> raw parameter array.
        """
        # Start with norm factor
        array_parts = [self.norm_factor]
        # Add weights w2 to wn
        array_parts.extend(self.weights)
        # Add sigmas s1 to sn
        array_parts.extend(self.sigmas)
        # Add background if requested
        if self.with_background and self.background is not None:
            array_parts.append(self.background)
        return np.array(array_parts)


##xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx shortcut functions
def cvt_params_to_stctd(params: np.ndarray, ngauss: int, with_background: bool = False) -> GaussParams:
    """
    raw parameter array -> structured GaussParams object.
    """
    return GaussParams.from_array(params, ngauss, with_background)


##xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx evaluate multi_gauss_model
class MGaussIn(BaseModel):
    crds: pnd.Np1DArray = Field(...)  ## coordinates, 用于 multi-gauss 计算的输入
    gpars: GaussParams  # MGauss 模型的参数


##xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx fitting input
# @dataclass
# class FitData:
#     """
#     Data model for fitting inputs: coordinates and values.
#     """

#     coords: pnd.Np1DArray
#     values: pnd.Np1DArray

#     @property
#     def x(self) -> pnd.Np1DArray:
#         """Alias for coordinates"""
#         return self.coords

#     @property
#     def y(self) -> pnd.Np1DArray:
#         """Alias for values"""
#         return self.values


class FitConfig(BaseModel):
    """
    Configuration parameters for multi-Gaussian fitting.
    """

    ngauss: int = 5
    with_background: bool = False
    is_post_bragg: bool = False
    lam_n: float = DftValues.penalty_norm
    info_verbose: bool = False
    info_tag: Optional[str] = None


##============================================================
## scipy.optimize.minimize, fun 作为要最小化的目标函数
## 函数原型为 fun(x, *args) -> float
## 以下class 打包 *args
class ObjctvArgs2(BaseModel):
    ## params: GaussParams
    crds: pnd.Np1DArray = Field(...)  ## coordinates, 用于 multi-gauss 计算的输入
    vals_tgt: pnd.Np1DArray = Field(...)  # 期望的 y 值
    fitcfg: FitConfig = Field(...)
    sigma_prev: Optional[pnd.Np1DArray] = None  ## sigmas previous
    mt_wts_pars: Optional[MidTailDcttRet] = None  ## mid tail weights parameters


##xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
class BoundsIn(BaseModel):
    """
    Parameters for optimization bounds and constraints.
    """

    ngauss: int
    crds: pnd.Np1DArray = Field(...)
    is_post_bragg: bool
    step_range: float = DftValues.step_range
    step_range_relax: float = DftValues.step_rlxd_range
    sig_prev: Optional[pnd.Np1DArray] = None  # sigmas previous


class CnstrntsIn(BaseModel):
    """
    计算约束输入
    """

    ngauss: int
    is_post_bragg: bool


class RunMGFitIn(BaseModel):
    crds: pnd.Np1DArray = Field(...)  ## coordinates, 用于 multi-gauss 计算的输入
    vals_tgt: pnd.Np1DArray = Field(...)  ## values target; 期望的 y 值
    fitcfg: Optional[FitConfig] = None
    sigma_prev: Optional[pnd.Np1DArray] = None
    gpars_init: Optional[GaussParams] = None


@dataclass
class RunMGFitRet:
    gpars0: GaussParams
    gpars1: GaussParams
    gpars2: GaussParams
    ret: OptimizeResult
